dict_filename <- 'Technical Test Data and Information/PUMS_Data_Dictionary_2019.csv'

dictionaries <- readr::read_lines(dict_filename) %>% 
  unique() %>% 
  split(str_detect(., '^NAME')) %>% # identify each dataset
  purrr::set_names(c('values', 'names')) %>% 
  purrr::map2(.y = list(val_columns = c('field_type',
                                        'column_name',
                                        'column_type',
                                        'field_length',
                                        'starting_value',
                                        'ending_value',
                                        'description'),
                        name_columns = c('field_type',
                                         'column_name',
                                         'column_type',
                                         'field_length',
                                         'description')),
              .f = ~data.table::fread(text = c(paste(.y, collapse = ','),
                                               .x)))

# find the index of SPORDER in both the $names and $values tables
sporder_index <- dictionaries %>% 
  purrr::map(~which(.x$column_name %chin% 'SPORDER'))

# map the corresponding dataset flag to each row
walk2(.x = dictionaries, 
      .y = sporder_index,
      .f = ~.x[, record_number := 1:.N
             ][, dataset := data.table::fifelse(test = record_number < .y,
                                           yes = 'H',
                                           no = 'P')]) 

dict_filename <- 'Technical Test Data and Information/PUMS_Data_Dictionary_2019.csv'
dicts <- parse_dictionaries(dict_filename)

str(object = dicts, give.atr = F, level = 2, vec.len = 2)

# load person data
per <- data.table::fread('Technical Test Data and Information/psam_p48.csv') %>% 
  discard_cols(cols = 'RT')
# 272,776 rows

# load housing unit data
hu <- data.table::fread('Technical Test Data and Information/psam_h48.csv') %>% 
  discard_cols(cols = 'RT')
# 123,710 rows

join_cols <- intersect(colnames(per), colnames(hu))
# this is c("SERIALNO", "DIVISION", "PUMA", "REGION", "ST", "ADJINC")

# set data.table keys for fast joining
data.table::setkeyv(per, join_cols)
data.table::setkeyv(hu,  join_cols)

# data.table left join
dat <- data.table::merge.data.table(x = per, hu, all.x = TRUE)

# remove the separate tables from memory
rm(per)
rm(hu)

# map all character variables as.character:
# first, get all the column names in the dictionary that correspond to 'C' column types. 
character_cols <- dictionaries$names[column_type == 'C']$column_name %>% intersect(colnames(dat))

# and then coerce all of them to character by reference
dat[, (character_cols) := purrr::map(.SD, as.character), .SDcols = character_cols]

dat <- discard_cols(dt = dat, cols = c('^WGTP', '^PWGTP'), regex = TRUE)

flag_columns <- dictionaries$names[str_detect(description, 'allocation flag')]$column_name
dat <- discard_cols(dt = dat, cols = flag_columns, regex = TRUE)

unique_values <- data.table(column = colnames(dat),
                            unique_values = purrr::map_int(dat, uniqueN)
                            )[unique_values == 1
                            ][, description := describe_column(column)][]

knitr::kable(unique_values)

dat <- discard_cols(dt = dat, cols = unique_values$column)

na_perc <- na_percentage(dat) %>% 
  data.table::merge.data.table(y = dictionaries$names[, .(column_name, column_type, description)],
                               by.x = 'column',
                               by.y = 'column_name')

na_perc[order(-na_perc)] %>% head(15) %>% knitr::kable()

numeric_impute <- dictionaries$names[column_type == 'N']$column_name %>%
  intersect(colnames(dat)) %>% 
  setdiff('PINCP') # our target variable deserves special consideration 

na_fill(dt = dat,  
        cols = numeric_impute,
        replacement = 0)

# for categoricals, only the variables that have an explicit 'b' = 'Does not apply' field will be replaced.
categorical_impute <- dictionaries$names[column_type == 'C']$column_name %>%
  intersect(colnames(dat)) %>% 
  intersect(dictionaries$values[column_type == 'C' & starting_value == 'b']$column_name)

na_fill(dt = dat,  
        cols = categorical_impute,
        replacement = 'b')

na_perc <- na_percentage(dat) %>% 
  data.table::merge.data.table(y = dictionaries$names[, .(column_name, column_type, description)],
                               by.x = 'column',
                               by.y = 'column_name')

na_perc[na_perc > 0] %>% head(10) %>% knitr::kable()

# select the 'C' variables according to the dictionary
columns_to_decode <- dictionaries$names[column_type %chin% 'C']$column_name 

# do not decode range categorical variables (ie where starting_value != ending_value)
# These are c("SERIALNO", "SPORDER", "PUMA", "MIGPUMA", "POWPUMA", "RACNUM")
exclude_from_decoding <- 
  dictionaries$values[column_type %chin% 'C' & starting_value != ending_value]$column_name

columns_to_decode %<>% setdiff(exclude_from_decoding) 

# and of course, decode only the columns that are still present in the data
columns_to_decode %<>% intersect(colnames(dat))

# and define a function that decodes a column
decode_categoricals <- function(col, dat){
  # select the dictionary entries speciying column values
  column_mapping <- dictionaries$values[column_name %chin% col & starting_value == ending_value,
                                        .(starting_value, description)] %>% 
    data.table::setkey(starting_value)
  
  # data.table 'left join' by reference
  # technically this is a mutate by reference, not a join
  matching_statement <- paste0(col, '==', 'starting_value')
  setkeyv(dat, col)
  dat[column_mapping, on = matching_statement, (col) := i.description]
}

tictoc::tic() 
# apply the function to all categorical columns
columns_to_decode %>% walk(decode_categoricals, dat = dat)
tictoc::toc()

per_fn <- 'Technical Test Data and Information/psam_p48.csv'
hu_fc <- 'Technical Test Data and Information/psam_h48.csv'
dict_filename <- 'Technical Test Data and Information/PUMS_Data_Dictionary_2019.csv'

pums <- pums_for_individual_modelling(person_filename = per_fn,
                                      housing_unit_filename = hu_fc,
                                      dictionary_filename = dict_filename)

num_summary(dt = dat, col = 'PINCP') %>% knitr::kable()

dat <- dat[!is.na(PINCP)]

chronicle::make_density(dt = dat, value = 'PINCP', static = FALSE)

# select numeric columns only
nums <- dat %>% purrr::keep(is.numeric) %>% colnames()

# compute correlations against the individual ncome variable, and count NA values 
income_cor <- data.table(column_name = nums,
                  income_cor = purrr::map_dbl(.x = nums,
                                              .f = ~cor(na.omit(dat[, .(PINCP, get(.x))]))[1,2] %>% round(2)),
                  na_perc = purrr::map_dbl(.x = nums, 
                                           .f = ~mean(is.na(dat[[.x]])))
                  )[, description := describe_column(column_name)]

income_cor[order(-abs(income_cor))] %>% knitr::kable()

dat <- discard_cols(dt = dat, cols = c('PERNP', 'WAGP', 'OIP', 
                                       'PAP',  'RETP', 'INTP', 
                                       'SEMP', 'SSIP', 'SSP', 
                                       'WAGP', 'WKHP'))

# create a copy of dat with a snippet of the variable description pasted into each column name
dat_with_descriptions <- data.table::copy(dat)
desc_index <- match(colnames(dat), dictionaries$names$column_name)
desc_colnames <- paste0(colnames(dat), ': ', str_sub(dictionaries$names$description[desc_index],1,30))
colnames(dat_with_descriptions) <- desc_colnames

skimr::skim(keep(dat_with_descriptions, is.numeric))

cor_dat <- cor(keep(dat, is.numeric)) %>% 
  round(2) %>% 
  as.data.table(keep.rownames = TRUE) %>% 
  data.table::melt.data.table(id.vars = 'rn', measure.vars = colnames(.)[-1]) %>% 
  setnames(old = c('rn', 'variable'), new = c('x', 'y'))

cor_dat[, x := factor(x, levels = sort(unique(x)))
      ][, y := factor(y, levels = sort(unique(x), decreasing = TRUE))]


(ggplot(data = cor_dat, 
                 mapping = aes(x = x, y = y, fill = value)) +
    geom_tile() +
    theme_minimal() +
    theme(panel.background = element_rect(fill = "transparent", colour = NA),
                   plot.background =  element_rect(fill = "transparent", colour = NA), 
                   axis.title = element_blank(),
                   axis.text.x = element_blank()) +
    viridis::scale_fill_viridis(option = 'magma')) %>% 
  plotly::ggplotly()

categorical_plots <- categoricals_to_explore %>% 
  purrr::map2(.x = categoricals_to_explore, 
              .y = describe_column(categoricals_to_explore), 
              .f = ~chronicle::make_boxplot(dt = dat[PINCP < 150000],
                                            value = 'PINCP', 
                                            groups = .x,
                                            jitter = FALSE, 
                                            x_axis_label = .y, 
                                            static = FALSE)) %>%
  set_names(categoricals_to_explore)

library(h2o)

#set up the H2O local cluster
h2o.init(max_mem_size = '16G') # allocating 16GB of RAM

h2o.no_progress() # disable progress bars

# feed all variables to the model (except indexes and the target variable)
x <- setdiff(colnames(dat), c('PINCP', 'SERIALNO', 'SPORDER'))
y <- 'PINCP'

# import dataset into h2o cluster
model_dat <- as.h2o(dat)

dat_split <- h2o.splitFrame(data = model_dat, ratios = 0.8, seed = 1)
train <- dat_split[[1]]
test <- dat_split[[2]]

# build a cross validated model
rf <- h2o.randomForest(x = x,
                       y = y,
                       training_frame = train,
                       nfolds = 10, 
                       max_runtime_secs = 100, 
                       seed = 1)

rf@model$model_summary %>% knitr::kable()

rf@model$cross_validation_metrics_summary %>% knitr::kable()

# check performance over test set
rf@model$cross_validation_metrics

h2o.performance(model = rf, newdata = test)

rf <- h2o.randomForest(x = x,
                       y = y,
                       training_frame = train,
                       nfolds = 10, 
                       max_runtime_secs = 600, 
                       seed = 1, 
                       max_depth = 30)

rf@model$model_summary %>% knitr::kable()

# check performance over test set
rf@model$cross_validation_metrics

h2o.performance(model = rf, newdata = test)

rf@model$scoring_history %>% knitr::kable()

aml <- h2o.automl(x = x,
                  y = y,
                  training_frame = train,
                  nfolds = 10, 
                  max_runtime_secs = 600, 
                  seed = 1)

aml@leaderboard %>% knitr::kable()

aml@leader@model$model_summary

# predict over the test set
pred <- h2o.predict(object = rf, 
                    newdata = test)

# compute relative error for each observation
pred_frame <- cbind(as.data.table(test)[, .(PINCP)],
                    as.data.table(pred)
                    )[, absolute_error := abs(predict-PINCP)
                    ][, relative_error := absolute_error/PINCP]

# create a grouping column for income quintiles
discretize_column(dt = pred_frame,
                  amount_column = 'PINCP',
                  band_cuts = quantile(pred_frame$PINCP, probs = seq(0, 1, .2)), 
                  bands_column_name = 'income_band')

# check the performance in each quintile
pred_frame[, .(population_size = .N,
               median_absolute_error = median(absolute_error),
               median_relative_error = median(relative_error)),
           by = income_band][order(income_band)] %>% knitr::kable()

# flag good predictions
pred_frame[, good_prediction := relative_error < .15]

pred_frame[, .(perc_good_pred = mean(good_prediction)), by = income_band][order(income_band)] %>% knitr::kable()

var_importance <- as.data.table(h2o.varimp(rf))[, cumulative_perc := cumsum(percentage)
                                              ][, relative_importance := NULL
                                              ][, description := describe_column(variable)]

head(var_importance, 10) %>% knitr::kable()

# fetch the column description from the NAME entries of the data dictionary
describe_column <- function(col = NULL){
  descrip_index <- match(col, dictionaries$names$column_name)
  descr <- dictionaries$names$description[descrip_index]
  
  if(any(is.na(descr))){
    descr[is.na(descr)] <- 'Not on ACS dictionary'
  }
  return(descr)
}

# run summary statistics over a numerical column of a data.table (optionally grouped)
num_summary <- function(dt, col, by = NULL){
  # rename the column to a generic variable (to avoid eval() or get() calls)
  data.table::setnames(x = dt, old = col, new = 'value')
  
  # perform the computations
  summ <- dt[, .(count = sum(!is.na(value)),
               distinct_values = uniqueN(value),
               na_count = sum(is.na(value)),
               na_perc = sum(is.na(value))/length(value),
               sum = sum(value, na.rm = TRUE),
               mean = mean(value, na.rm = TRUE),
               sd = sd(value, na.rm = TRUE),
               min = min(value, na.rm = TRUE),
               percentile.25 = quantile(value, 0.25, na.rm = TRUE),
               percentile.50 = quantile(value, 0.5, na.rm = TRUE),
               percentile.75 = quantile(value, 0.75, na.rm = TRUE),
               percentile.975 = quantile(value, 0.975, na.rm = TRUE),
               percentile.99 = quantile(value, 0.99, na.rm = TRUE),
               max = max(value, na.rm = TRUE)),
             by = by]
  
  # rename the column back to its original name
  data.table::setnames(x = dt, old = 'value', new = col)
  
  return(summ)
}

# compute the count and percentage of missing entries for each column of a data.frame 
na_percentage <- function(dt){
  data.table(column = colnames(dt),
             na_count = purrr::map_int(.x = dt, 
                                       .f = ~sum(is.na(.x))
                                       ))[, na_perc := na_count/nrow(dt)][]
}

# removes all columns matching a character vector (optionally regular expressions)
discard_cols <- function(dt, cols, regex = FALSE){
  if(!regex){
    dt %>% purrr::discard(colnames(.) %in% cols)
  }else{
    dt %>% purrr::discard(str_detect(colnames(.), paste(cols, collapse = '|')))
  }
}


na_fill = function(dt, cols, replacement = 0) {
  # match the indexes of the columns specified by the user
  index <- match(cols, colnames(dt))
  purrr::walk(index, ~data.table::set(dt, which(is.na(dt[[.x]])), .x, replacement))
}


# parse the two separate specifications entangled in the ACS dictionary file
parse_dictionaries <- function(dict_filename){
  dictionaries <- readr::read_lines(dict_filename) %>%
    unique() %>%
    split(stringr::str_detect(., '^NAME')) %>% # identify each dataset
    purrr::set_names(c('values', 'names')) %>%
    purrr::map2(.y = list(val_columns = c('field_type',
                                          'column_name',
                                          'column_type',
                                          'field_length',
                                          'starting_value',
                                          'ending_value',
                                          'description'),
                          name_columns = c('field_type',
                                           'column_name',
                                           'column_type',
                                           'field_length',
                                           'description')),
                .f = ~data.table::fread(text = c(paste(.y, collapse = ','),
                                                 .x)))
  
  # find the index of SPORDER in both the $names and $values tables
  sporder_index <- dictionaries %>%
    purrr::map(~which(.x$column_name %chin% 'SPORDER'))
  
  # map the corresponding dataset flag to each row
  walk2(.x = dictionaries,
        .y = sporder_index,
        .f = ~.x[, record_number := 1:.N
        ][, dataset := data.table::fifelse(test = record_number < .y,
                                           yes = 'H',
                                           no = 'P')])
  return(dictionaries)
}

# decode categorical columns with the values provided in the VAL dictionary.
decode_categoricals <- function(col, dat){
  # select the dictionary entries speciying column values
  column_mapping <- dictionaries$values[column_name %chin% col & starting_value == ending_value,
                                        .(starting_value, description)] %>% 
    data.table::setkey(starting_value)
  
  # data.table 'left join' by reference
  # technically this is a mutate by reference, not a join
  matching_statement <- paste0(col, '==', 'starting_value')
  setkeyv(dat, col)
  dat[column_mapping, on = matching_statement, (col) := i.description]
}

pums_for_individual_modelling <- function(person_filename,
                                          housing_unit_filename,
                                          dictionary_filename){

  tictoc::tic('PUMS files successfully parsed')
  # 0. load and join data ---
    
  # load person data
  per <- data.table::fread('Technical Test Data and Information/psam_p48.csv') %>% 
    discard_cols(cols = 'RT')

  # load housing unit data
  hu <- data.table::fread('Technical Test Data and Information/psam_h48.csv') %>% 
    discard_cols(cols = 'RT')

  dictionaries <- parse_dictionaries(dict_filename)
  
  join_cols <- intersect(colnames(per), colnames(hu))
  # this is c("SERIALNO", "DIVISION", "PUMA", "REGION", "ST", "ADJINC")
  
  # set data.table keys for fast joining
  data.table::setkeyv(per, join_cols)
  data.table::setkeyv(hu,  join_cols)
  
  # data.table left join
  pums <- data.table::merge.data.table(x = per, hu, all.x = TRUE)
  
  character_cols <- dictionaries$names[column_type == 'C']$column_name %>% intersect(colnames(dat))
  pums[, (character_cols) := purrr::map(.SD, as.character), .SDcols = character_cols]
  
  # remove the separate tables from memory
  rm(per)
  rm(hu)
  
  # 1. discard columns ----
  
  # weight columns
  pums <- discard_cols(dt = pums, cols = c('^WGTP', '^PWGTP'), regex = TRUE)
  
  # flag columns
  flag_columns <- dictionaries$names[str_detect(description, 'allocation flag')]
  pums <- discard_cols(dt = pums, cols = flag_columns, regex = TRUE)
  
  # constant_columns
  unique_values <- data.table(column = colnames(pums),
                              unique_values = purrr::map_int(pums, uniqueN)
                              )[unique_values == 1]
  
  pums <- discard_cols(dt = pums, cols = unique_values$column)
  
  # 2. data imputation ----
  numeric_impute <- dictionaries$names[column_type == 'N']$column_name %>%
    intersect(colnames(pums)) %>% 
    setdiff('PINCP') # our target variable deserves special consideration 
  
  na_fill(dt = pums,  
          cols = numeric_impute,
          replacement = 0)
  
  # for categoricals, only the variables that have an explicit 'b' = 'Does not apply' field will be replaced.
  categorical_impute <- dictionaries$names[column_type == 'C']$column_name %>%
    intersect(colnames(pums)) %>% 
    intersect(dictionaries$values[column_type == 'C' & starting_value == 'b']$column_name)
  
  na_fill(dt = pums,  
          cols = categorical_impute,
          replacement = 'b')
  
  
  # 3. categorical variable decoding ----
  
  # select the categorical variables according to the dictionary
  columns_to_decode <- dictionaries$names[column_type %chin% 'C']$column_name 
  
  # do not decode range categorical variables (ie where starting_value != ending_value)
  # These are c("SERIALNO", "SERIALNO", "PUMA", "MIGPUMA", "POWPUMA", "RACNUM")
  exclude_from_decoding <- 
    dictionaries$values[column_type %chin% 'C' & starting_value != ending_value]$column_name
  
  columns_to_decode %<>% setdiff(exclude_from_decoding) 
  
  # and remove the columns we already removed from the dataaset
  columns_to_decode %<>% intersect(colnames(pums))
  
  # set all categorical variables as character vectors to avoid type inconsistencies 
  pums[, (columns_to_decode) := purrr::map(.SD, as.character), .SDcols = columns_to_decode]
  
  columns_to_decode %>% walk(decode_categoricals, dat = pums)
  
  tictoc::toc()
  
  return(pums[])
}

# create a new column marking the value band of each observation wrt a continuous column 
discretize_column <- function(dt,
                              amount_column,
                              band_cuts,
                              bands_column_name) {
  
  amount_column %<>% as.character()
  
  # create empty column
  data.table::setDT(dt)
  dt[, (bands_column_name) := NA_character_]
  
  # build index for the bands
  run <- seq_along(band_cuts)
  band_length <- length(band_cuts)
  
  # create names for the bands
  band_names <- paste0(band_cuts[-band_length], '-', band_cuts[-1])
  
  # assign band to each value
  purrr::walk(run, ~data.table::set(dt,
                                    which(band_cuts[.x] <= dt[[amount_column]] &
                                            dt[[amount_column]] <= band_cuts[.x + 1]),
                                    bands_column_name,
                                    band_names[.x])[])
  
  dt[, (bands_column_name) := factor(x = get(bands_column_name), levels = band_names)]
}

sessionInfo()